"""
python fujian2_lstm_predict.py
"""

import os
import json
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import tensorflow as tf
import re

def read_json_data(file_path):
    with open(file_path, 'r') as f:
        data = json.load(f)
    df = pd.DataFrame(data)
    df['date'] = pd.to_datetime(df['date'])
    df.rename(columns={'sales': 'y'}, inplace=True)
    return df

def prepare_data(df):
    values = df['y'].values
    input_length = len(values)
    sequence_length = min(50, input_length)
    X = []
    y = []
    for i in range(input_length - sequence_length):
        X.append(values[i:i + sequence_length])
        y.append(values[i + sequence_length])
    X = np.array(X)
    y = np.array(y)
    X = X.reshape(X.shape[0], sequence_length, 1)
    return X, y

def train_lstm(X_train, y_train):
    input_layer = tf.keras.Input(shape=(X_train.shape[1], X_train.shape[2]))
    lstm_layer = tf.keras.layers.LSTM(100)(input_layer)
    output_layer = tf.keras.layers.Dense(1)(lstm_layer)
    model = tf.keras.Model(inputs=input_layer, outputs=output_layer)
    model.compile(loss='mae', optimizer='adam')
    model.fit(X_train, y_train, epochs=100, batch_size=32, verbose=0)
    return model

def predict_future_sales(model, last_sequence, num_days):
    predictions = []
    for _ in range(num_days):
        new_prediction = model.predict(last_sequence.reshape(1, last_sequence.shape[0], 1))[0][0]
        predictions.append(new_prediction)
        last_sequence = np.append(last_sequence[1:], new_prediction)
    return predictions

def save_predictions_to_json(category_id, predictions, output_folder):
    future_dates = pd.date_range(start=pd.to_datetime('2023-07-01'), end=pd.to_datetime('2023-09-30'))
    result = [{'date': date.strftime('%Y/%m/%d'), 'predicted_sales': int(sale)} for date, sale in zip(future_dates, predictions)]
    output_json_path = os.path.join(output_folder, f'json_category{category_id}.json')
    os.makedirs(os.path.dirname(output_json_path), exist_ok=True)
    with open(output_json_path, 'w') as f:
        json.dump(result, f)

def plot_data_and_predictions(df, predictions, category_id, output_folder):
    original_dates = df['date']
    original_sales = df['y']
    future_dates = pd.date_range(start=pd.to_datetime('2023-07-01'), end=pd.to_datetime('2023-09-30'))
    plt.figure(figsize=(10, 6))
    plt.scatter(original_dates, original_sales, label='Original Data', c='blue')
    plt.scatter(future_dates, predictions, label='Predicted Data', c='red')
    plt.legend()
    plt.xlabel('Date')
    plt.ylabel('Sales')
    plt.title(f'LSTM Prediction on Category {category_id}')
    output_image_path = os.path.join(output_folder, f'gragh_category{category_id}.png')
    os.makedirs(os.path.dirname(output_image_path), exist_ok=True)
    plt.savefig(output_image_path)

if __name__ == "__main__":
    input_folder = r'..\fujian\fujian2\groupByCategory'
    output_folder = r'./test'
    json_files = [os.path.join(input_folder, file) for file in os.listdir(input_folder) if file.endswith('.json')]
    for json_file in json_files:
        df = read_json_data(json_file)
        X, y = prepare_data(df)
        split_index = int(0.8 * len(df))
        if split_index > 0:
            X_train, X_test = X[:split_index], X[split_index:]
            y_train, y_test = y[:split_index], y[split_index:]
            model = train_lstm(X_train, y_train)
            last_sequence = X_test[-1] if len(X_test) > 0 else None
            filename = os.path.basename(json_file)
            category_id = -1
            match = re.search(r'category_(?:category)?(\d+)', filename)
            if match:
                category_id = int(match.group(1))
            else:
                print("未找到匹配的 category ID")
            if last_sequence is not None:
                future_predictions = predict_future_sales(model, last_sequence, 92)
                save_predictions_to_json(category_id, future_predictions, output_folder)
                plot_data_and_predictions(df, future_predictions, category_id, output_folder)
